#pragma once

#include "function.h"

namespace  ldl
{
namespace autograd
{
template<typename T>
class Sub : public Function<T>
{
public:
    Sub() {
        Function<T>::m_inputs.resize(2);
        Function<T>::m_outputs.resize(1);
        Function<T>::m_outputs[0] = std::make_shared<Autograd<T>>();
    }
    virtual ~Sub() = default;

    virtual void forward() override {
        auto& data_in0 = *(Function<T>::m_inputs[0]->data());
        auto& data_in1 = *(Function<T>::m_inputs[1]->data());
        auto& data_ptr = Function<T>::m_outputs_weak_ptr[0].lock()->data();

        if(data_ptr != nullptr)[[likely]] {
            *data_ptr = data_in0 - data_in1;
        } else[[unlikely]] {
            data_ptr = std::make_shared<T>(data_in0 - data_in1);
        }
    }
    virtual void do_backward() override {
        auto& grad_out = *(Function<T>::m_outputs_weak_ptr[0].lock()->grad());
        auto& grad_in0 = *(Function<T>::m_inputs[0]->grad());
        auto& grad_in1 = *(Function<T>::m_inputs[1]->grad());

        grad_in0 += grad_out;
        grad_in1 -= grad_out;
    }
};
}
}
